// Copyright (c) 2005-2008, Simon Howard
// Permission to use, copy, modify, and/or distribute this software
// for any purpose with or without fee is hereby granted, provided
// that the above copyright notice and this permission notice appear
// in all copies.
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL
// WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE
// AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT, OR
// CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
// NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
// CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.

#![allow(dead_code)]

use std::cmp::Ordering;
use std::mem;

const NUM_TEST_VALUES: usize = 1000;

#[derive(Debug)]
struct AVLTree<K, V> {
    root: Option<Box<AVLTreeNode<K, V>>>,
    num_entries: usize,
}

#[derive(Debug)]
struct AVLTreeNode<K, V> {
    key: K,
    value: V,
    height: i32,
    left: Option<Box<AVLTreeNode<K, V>>>,
    right: Option<Box<AVLTreeNode<K, V>>>,
}

impl<K, V> AVLTree<K, V>
where
    K: Ord,
{
    /// Creates a new empty AVL tree.
    pub fn new() -> Self {
        AVLTree {
            root: None,
            num_entries: 0,
        }
    }

    /// Inserts a key-value pair into the tree.
    pub fn insert(&mut self, key: K, value: V) {
        self.root = Self::insert_node(self.root.take(), key, value);
        self.num_entries += 1;
    }

    fn insert_node(
        node: Option<Box<AVLTreeNode<K, V>>>,
        key: K,
        value: V,
    ) -> Option<Box<AVLTreeNode<K, V>>> {
        if let Some(mut node) = node {
            match key.cmp(&node.key) {
                Ordering::Less => {
                    node.left = Self::insert_node(node.left.take(), key, value);
                }
                Ordering::Greater => {
                    node.right = Self::insert_node(node.right.take(), key, value);
                }
                Ordering::Equal => {
                    node.value = value;
                }
            }
            node.height = 1 + std::cmp::max(Self::node_height(&node.left), Self::node_height(&node.right));
            Some(Self::balance(node))
        } else {
            Some(Box::new(AVLTreeNode {
                key,
                value,
                height: 1,
                left: None,
                right: None,
            }))
        }
    }

    fn balance(mut node: Box<AVLTreeNode<K, V>>) -> Box<AVLTreeNode<K, V>> {
        let balance_factor = Self::node_height(&node.left) - Self::node_height(&node.right);
        if balance_factor > 1 {
            if Self::node_height(&node.left.as_ref().unwrap().left) < Self::node_height(&node.left.as_ref().unwrap().right) {
                node.left = Self::rotate_left(node.left.take().unwrap());
            }
            return Self::rotate_right(node);
        } else if balance_factor < -1 {
            if Self::node_height(&node.right.as_ref().unwrap().right) < Self::node_height(&node.right.as_ref().unwrap().left) {
                node.right = Self::rotate_right(node.right.take().unwrap());
            }
            return Self::rotate_left(node);
        }
        node
    }

    fn rotate_left(mut node: Box<AVLTreeNode<K, V>>) -> Box<AVLTreeNode<K, V>> {
        let mut new_root = node.right.take().unwrap();
        node.right = new_root.left.take();
        new_root.left = Some(node);
        new_root.height = 1 + std::cmp::max(Self::node_height(&new_root.left), Self::node_height(&new_root.right));
        new_root
    }

    fn rotate_right(mut node: Box<AVLTreeNode<K, V>>) -> Box<AVLTreeNode<K, V>> {
        let mut new_root = node.left.take().unwrap();
        node.left = new_root.right.take();
        new_root.right = Some(node);
        new_root.height = 1 + std::cmp::max(Self::node_height(&new_root.left), Self::node_height(&new_root.right));
        new_root
    }

    fn node_height(node: &Option<Box<AVLTreeNode<K, V>>>) -> i32 {
        node.as_ref().map_or(0, |n| n.height)
    }

    /// Looks up a value by key.
    pub fn lookup(&self, key: &K) -> Option<&V> {
        let mut current = &self.root;
        while let Some(node) = current {
            match key.cmp(&node.key) {
                Ordering::Less => current = &node.left,
                Ordering::Greater => current = &node.right,
                Ordering::Equal => return Some(&node.value),
            }
        }
        None
    }

    /// Removes a key-value pair from the tree.
    pub fn remove(&mut self, key: &K) -> Option<V> {
        let (root, removed_value) = Self::remove_node(self.root.take(), key);
        self.root = root;
        if removed_value.is_some() {
            self.num_entries -= 1;
        }
        removed_value
    }

    fn remove_node(
        node: Option<Box<AVLTreeNode<K, V>>>,
        key: &K,
    ) -> (Option<Box<AVLTreeNode<K, V>>>, Option<V>) {
        if let Some(mut node) = node {
            match key.cmp(&node.key) {
                Ordering::Less => {
                    let (left, removed_value) = Self::remove_node(node.left.take(), key);
                    node.left = left;
                    (Some(Self::balance(node)), removed_value)
                }
                Ordering::Greater => {
                    let (right, removed_value) = Self::remove_node(node.right.take(), key);
                    node.right = right;
                    (Some(Self::balance(node)), removed_value)
                }
                Ordering::Equal => {
                    let value = mem::replace(&mut node.value, unsafe { mem::zeroed() });
                    let new_root = Self::merge(node.left.take(), node.right.take());
                    (new_root, Some(value))
                }
            }
        } else {
            (None, None)
        }
    }

    fn merge(
        left: Option<Box<AVLTreeNode<K, V>>>,
        right: Option<Box<AVLTreeNode<K, V>>>,
    ) -> Option<Box<AVLTreeNode<K, V>>> {
        if let Some(mut right_node) = right {
            if let Some(left_node) = left {
                let (new_left, _) = Self::remove_node(Some(left_node), &right_node.key);
                right_node.left = new_left;
                right_node.height = 1 + std::cmp::max(Self::node_height(&right_node.left), Self::node_height(&right_node.right));
                Some(Self::balance(right_node))
            } else {
                Some(right_node)
            }
        } else {
            left
        }
    }

    /// Returns the number of entries in the tree.
    pub fn num_entries(&self) -> usize {
        self.num_entries
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_avl_tree_new() {
        let tree: AVLTree<i32, i32> = AVLTree::new();
        assert_eq!(tree.num_entries(), 0);
        assert!(tree.root.is_none());
    }

    #[test]
    fn test_avl_tree_insert_lookup() {
        let mut tree = AVLTree::new();
        for i in 0..NUM_TEST_VALUES {
            tree.insert(i, i);
            assert_eq!(tree.num_entries(), i + 1);
        }
        assert!(tree.root.is_some());

        for i in 0..NUM_TEST_VALUES {
            assert_eq!(tree.lookup(&i), Some(&i));
        }

        assert_eq!(tree.lookup(&(NUM_TEST_VALUES + 100)), None);
    }

    #[test]
    fn test_avl_tree_remove() {
        let mut tree = AVLTree::new();
        for i in 0..NUM_TEST_VALUES {
            tree.insert(i, i);
        }

        for i in 0..NUM_TEST_VALUES {
            assert_eq!(tree.remove(&i), Some(i));
            assert_eq!(tree.num_entries(), NUM_TEST_VALUES - i - 1);
        }

        assert_eq!(tree.remove(&(NUM_TEST_VALUES + 100)), None);
        assert_eq!(tree.num_entries(), 0);
    }
}
